from fastapi import FastAPI, Depends, HTTPException, Query
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from pydantic import BaseModel, Field
from datetime import date, datetime, timedelta
from typing import List, Optional, Dict, Any
import jwt
from passlib.context import CryptContext
from jose import JWTError, jwt
from sqlalchemy import create_engine, Column, Integer, String, Date, Float, ForeignKey
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session, relationship
from enum import Enum
import random
from typing import Union

# 数据库配置
SQLALCHEMY_DATABASE_URL = "sqlite:///./attendance.db"
engine = create_engine(SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

# 安全配置
SECRET_KEY = "your-secret-key"
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30

# 密码哈希
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# OAuth2
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

# 部门枚举
class Department(str, Enum):
    HR = "人力资源部"
    TECH = "技术部"
    MARKETING = "市场部"
    FINANCE = "财务部"
    ADMIN = "行政部"
    SALES = "销售部"

# 考勤类型枚举
class AttendanceType(str, Enum):
    NORMAL = "正常出勤"
    LATE = "迟到"
    EARLY = "早退"
    ABSENT = "缺勤"
    LEAVE = "请假"

# 数据模型
class User(Base):
    __tablename__ = "users"
    
    id = Column(Integer, primary_key=True, index=True)
    username = Column(String, unique=True, index=True)
    hashed_password = Column(String)
    department = Column(String)
    role = Column(String)  # admin, manager, employee
    
    attendance_records = relationship("AttendanceRecord", back_populates="user")

class AttendanceRecord(Base):
    __tablename__ = "attendance_records"
    
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    date = Column(Date)
    check_in_time = Column(String)
    check_out_time = Column(String)
    status = Column(String)  # normal, late, absent, leave
    duration = Column(Float)  # 工作时长
    
    user = relationship("User", back_populates="attendance_records")

# 创建数据库表
Base.metadata.create_all(bind=engine)

# Pydantic模型
class UserCreate(BaseModel):
    username: String
    password: String
    department: Department
    role: String

class UserLogin(BaseModel):
    username: String
    password: String

class UserResponse(BaseModel):
    id: int
    username: String
    department: Department
    role: String
    
    class Config:
        orm_mode = True

class AttendanceRecordCreate(BaseModel):
    user_id: int
    date: date
    check_in_time: Optional[str] = None
    check_out_time: Optional[str] = None
    status: str
    duration: Optional[float] = None

class AttendanceRecordResponse(BaseModel):
    id: int
    user_id: int
    date: date
    check_in_time: Optional[str] = None
    check_out_time: Optional[str] = None
    status: str
    duration: Optional[float] = None
    
    class Config:
        orm_mode = True

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: Optional[str] = None

# 依赖项
def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()

# 验证密码
def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

# 哈希密码
def get_password_hash(password):
    return pwd_context.hash(password)

# 获取用户
def get_user(db: Session, username: str):
    return db.query(User).filter(User.username == username).first()

# 认证用户
def authenticate_user(db: Session, username: str, password: str):
    user = get_user(db, username)
    if not user:
        return False
    if not verify_password(password, user.hashed_password):
        return False
    return user

# 创建访问令牌
def create_access_token(data: dict, expires_delta: Optional[timedelta] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=15)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt

# 获取当前用户
async def get_current_user(token: str = Depends(oauth2_scheme), db: Session = Depends(get_db)):
    credentials_exception = HTTPException(
        status_code=401,
        detail="无法验证凭据",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception
    user = get_user(db, username=token_data.username)
    if user is None:
        raise credentials_exception
    return user

# 获取当前活动用户
async def get_current_active_user(current_user: User = Depends(get_current_user)):
    if current_user.disabled:
        raise HTTPException(status_code=400, detail="无效用户")
    return current_user

# 初始化数据库
def init_db():
    db = SessionLocal()
    
    # 创建管理员用户
    admin_user = db.query(User).filter(User.username == "admin").first()
    if not admin_user:
        admin_user = User(
            username="admin",
            hashed_password=get_password_hash("admin123"),
            department=Department.HR,
            role="admin"
        )
        db.add(admin_user)
        db.commit()
        db.refresh(admin_user)
    
    # 创建示例用户
    users_data = [
        {"username": "zhangsan", "department": Department.TECH, "role": "employee"},
        {"username": "lisi", "department": Department.MARKETING, "role": "employee"},
        {"username": "wangwu", "department": Department.FINANCE, "role": "employee"},
        {"username": "zhaoliu", "department": Department.HR, "role": "employee"},
        {"username": "manager1", "department": Department.TECH, "role": "manager"},
    ]
    
    for user_data in users_data:
        user = db.query(User).filter(User.username == user_data["username"]).first()
        if not user:
            user = User(
                username=user_data["username"],
                hashed_password=get_password_hash("password"),
                department=user_data["department"],
                role=user_data["role"]
            )
            db.add(user)
            db.commit()
            db.refresh(user)
    
    # 生成示例考勤数据
    today = date.today()
    for i in range(30):  # 生成30天的考勤数据
        current_date = today - timedelta(days=i)
        for user in db.query(User).all():
            # 随机生成考勤记录
            record = db.query(AttendanceRecord).filter(
                AttendanceRecord.user_id == user.id,
                AttendanceRecord.date == current_date
            ).first()
            
            if not record:
                # 周末休息
                if current_date.weekday() >= 5:
                    continue
                
                # 随机状态
                statuses = ["normal", "late", "absent", "leave"]
                weights = [0.75, 0.12, 0.03, 0.10]  # 正常、迟到、缺勤、请假的概率
                status = random.choices(statuses, weights)[0]
                
                # 计算工作时长
                if status == "normal":
                    check_in_time = f"{random.randint(8, 9)}:{random.randint(0, 30):02d}"
                    check_out_time = f"{random.randint(17, 18)}:{random.randint(0, 30):02d}"
                    duration = random.uniform(8, 9)
                elif status == "late":
                    check_in_time = f"{random.randint(9, 10)}:{random.randint(0, 30):02d}"
                    check_out_time = f"{random.randint(17, 18)}:{random.randint(0, 30):02d}"
                    duration = random.uniform(7, 8)
                else:
                    check_in_time = None
                    check_out_time = None
                    duration = 0
                
                record = AttendanceRecord(
                    user_id=user.id,
                    date=current_date,
                    check_in_time=check_in_time,
                    check_out_time=check_out_time,
                    status=status,
                    duration=duration
                )
                db.add(record)
                db.commit()
    
    db.close()

# 创建FastAPI应用
app = FastAPI(title="考勤系统API", description="用于管理和统计员工考勤的后端API")

# 初始化数据库
init_db()

# API路由
@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends(), db: Session = Depends(get_db)):
    user = authenticate_user(db, form_data.username, form_data.password)
    if not user:
        raise HTTPException(
            status_code=400,
            detail="用户名或密码不正确",
            headers={"WWW-Authenticate": "Bearer"},
        )
    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": user.username}, expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}

@app.get("/users/me/", response_model=UserResponse)
async def read_users_me(current_user: User = Depends(get_current_active_user)):
    return current_user

# 获取所有用户
@app.get("/users/", response_model=List<UserResponse>)
async def get_users(
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    # 只有管理员和经理可以查看所有用户
    if current_user.role not in ["admin", "manager"]:
        raise HTTPException(status_code=403, detail="权限不足")
    
    users = db.query(User).all()
    return users

# 获取用户考勤记录
@app.get("/attendance/user/{user_id}", response_model=List<AttendanceRecordResponse>)
async def get_user_attendance(
    user_id: int,
    start_date: Optional[date] = None,
    end_date: Optional[date] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    # 只有管理员、经理或用户本人可以查看考勤记录
    if current_user.role not in ["admin", "manager"] and current_user.id != user_id:
        raise HTTPException(status_code=403, detail="权限不足")
    
    query = db.query(AttendanceRecord).filter(AttendanceRecord.user_id == user_id)
    
    if start_date:
        query = query.filter(AttendanceRecord.date >= start_date)
    if end_date:
        query = query.filter(AttendanceRecord.date <= end_date)
    
    records = query.all()
    return records

# 获取部门考勤统计
@app.get("/statistics/department/")
async def get_department_statistics(
    department: Optional[Department] = None,
    start_date: Optional[date] = None,
    end_date: Optional[date] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    # 只有管理员和经理可以查看部门统计
    if current_user.role not in ["admin", "manager"]:
        raise HTTPException(status_code=403, detail="权限不足")
    
    # 查询条件
    query = db.query(AttendanceRecord)
    
    if department:
        users_in_department = db.query(User.id).filter(User.department == department).subquery()
        query = query.filter(AttendanceRecord.user_id.in_(users_in_department))
    
    if start_date:
        query = query.filter(AttendanceRecord.date >= start_date)
    if end_date:
        query = query.filter(AttendanceRecord.date <= end_date)
    
    # 计算统计数据
    total_records = query.count()
    normal_records = query.filter(AttendanceRecord.status == "normal").count()
    late_records = query.filter(AttendanceRecord.status == "late").count()
    absent_records = query.filter(AttendanceRecord.status == "absent").count()
    leave_records = query.filter(AttendanceRecord.status == "leave").count()
    
    attendance_rate = (normal_records / total_records) * 100 if total_records > 0 else 0
    
    return {
        "department": department,
        "start_date": start_date,
        "end_date": end_date,
        "total_records": total_records,
        "attendance_rate": attendance_rate,
        "late_records": late_records,
        "absent_records": absent_records,
        "leave_records": leave_records
    }

# 获取所有部门考勤统计
@app.get("/statistics/departments/")
async def get_all_departments_statistics(
    start_date: Optional[date] = None,
    end_date: Optional[date] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    # 只有管理员和经理可以查看部门统计
    if current_user.role not in ["admin", "manager"]:
        raise HTTPException(status_code=403, detail="权限不足")
    
    statistics = []
    
    for department in Department:
        stats = await get_department_statistics(
            department=department,
            start_date=start_date,
            end_date=end_date,
            db=db,
            current_user=current_user
        )
        statistics.append(stats)
    
    return statistics

# 获取月度考勤趋势
@app.get("/statistics/trend/monthly/")
async def get_monthly_attendance_trend(
    year: int = Query(..., description="年份"),
    month: int = Query(..., description="月份"),
    department: Optional[Department] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_active_user)
):
    # 只有管理员和经理可以查看统计数据
    if current_user.role not in ["admin", "manager"]:
        raise HTTPException(status_code=403, detail="权限不足")
    
    # 获取当月的天数
    from calendar import monthrange
    _, days_in_month = monthrange(year, month)
    
    # 准备日期列表
    dates = [date(year, month, day) for day in range(1, days_in_month + 1)]
    
    # 准备数据
    attendance_data = []
    late_data = []
    
    for day in dates:
        # 查询条件
        query = db.query(AttendanceRecord)
        
        if department:
            users_in_department = db.query(User.id).filter(User.department == department).subquery()
            query = query.filter(AttendanceRecord.user_id.in_(users_in_department))
        
        query = query.filter(AttendanceRecord.date == day)
        
        # 计算出勤人数和迟到人数
        total_users = len(db.query(User).filter(User.department == department).all()) if department else len(db.query(User).all())
        attendance_count = query.filter(AttendanceRecord.status == "normal").count()
        late_count = query.filter(AttendanceRecord.status == "late").count()
        
        # 计算出勤率和迟到率
        attendance_rate = (attendance_count / total_users) * 100 if total_users > 0 else 0
        
        attendance_data.append(attendance_rate)
        late_data.append(late_count)
    
    # 返回数据
    return {
        "labels": [d.strftime("%d日") for d in dates],
        "attendance_data": attendance_data,
        "late_data": late_data
    }

# 主函数
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.000.0", port=8000)  